package com.strp.mptcp.congestion;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 多路径拥塞控制器
 * 支持多种拥塞控制算法：
 * - AIMD (默认)
 * - BBR
 * - CUBIC
 *
 * 主要功能：
 * 1. 为每个子流维护独立的拥塞控制状态
 * 2. 支持动态切换拥塞控制算法
 * 3. 提供 RTT 和带宽估计
 * 4. 支持拥塞窗口和发送速率控制
 *
 * @author 曹行阳
 * @date 2025/6/12
 */
public class CongestionController {
    private final ConcurrentHashMap<Integer, PathCongestionInfo> pathStats = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<Integer, CongestionAlgorithm> algorithms = new ConcurrentHashMap<>();

    // 拥塞控制算法配置
    private static final double INITIAL_CWND = 10.0;       // 初始拥塞窗口 (包数)
    private static final double INITIAL_SSTHRESH = 65535;  // 初始慢启动阈值 (包数)
    private static final double ALPHA = 1.0 / 8;           // srtt 平滑系数
    private static final double BETA = 1.0 / 4;            // rttvar 平滑系数
    private static final double GAMMA = 0.1;               // BBR 探测系数
    private static final int MIN_RTT_WINDOW_MS = 10000;    // BBR 最小 RTT 窗口
    private static final int MAX_RTT_WINDOW_MS = 30000;    // BBR 最大 RTT 窗口

    public enum CongestionAlgorithm {
        AIMD,   // 加性增乘性减
        BBR,    // Google BBR
        CUBIC   // CUBIC
    }

    /**
     * 注册一条子流，用于后续拥塞控制
     */
    public void registerPath(int pathId, CongestionAlgorithm algorithm) {
        pathStats.putIfAbsent(pathId, new PathCongestionInfo(INITIAL_CWND, INITIAL_SSTHRESH));
        algorithms.put(pathId, algorithm);
    }

    /**
     * 每当一个数据包被发送时，调用此方法
     */
    public void onSend(int pathId, int bytesSent) {
        PathCongestionInfo info = pathStats.get(pathId);
        if (info != null) {
            info.addBytesSent(bytesSent);
            info.packetSent();

            // 更新 BBR 状态
            if (algorithms.get(pathId) == CongestionAlgorithm.BBR) {
                updateBBRState(info, bytesSent);
            }
        }
    }

    /**
     * 收到 ACK 时调用
     */
    public void onAck(int pathId, int bytesAcked) {
        PathCongestionInfo info = pathStats.get(pathId);
        if (info == null) return;

        info.packetAcked();
        info.addBytesAcked(bytesAcked);

        CongestionAlgorithm algorithm = algorithms.get(pathId);
        switch (algorithm) {
            case AIMD -> updateAIMD(info);
            case BBR -> updateBBR(info);
            case CUBIC -> updateCUBIC(info);
        }
    }

    /**
     * 检测到丢包时调用
     */
    public void onLoss(int pathId) {
        PathCongestionInfo info = pathStats.get(pathId);
        if (info == null) return;

        CongestionAlgorithm algorithm = algorithms.get(pathId);
        switch (algorithm) {
            case AIMD -> handleAIMDLoss(info);
            case BBR -> handleBBRLoss(info);
            case CUBIC -> handleCUBICLoss(info);
        }
    }

    /**
     * 检查是否可以发送数据
     */
    public boolean canSend(int pathId) {
        PathCongestionInfo info = pathStats.get(pathId);
        if (info == null) return false;

        CongestionAlgorithm algorithm = algorithms.get(pathId);
        return switch (algorithm) {
            case AIMD -> info.getInflight() < info.cwnd;
            case BBR -> info.getInflight() < info.bbrCwnd;
            case CUBIC -> info.getInflight() < info.cubicCwnd;
        };
    }

    private void updateAIMD(PathCongestionInfo info) {
        if (info.cwnd < info.ssthresh) {
            // 慢启动阶段：指数增长
            info.cwnd = Math.min(info.cwnd * 2, info.ssthresh);
        } else {
            // 拥塞避免阶段：线性增长
            info.cwnd += 1.0 / info.cwnd;
        }
    }

    private void handleAIMDLoss(PathCongestionInfo info) {
        info.ssthresh = Math.max(info.cwnd / 2, 2.0);
        info.cwnd = 1.0;
        info.resetInflight();
    }

    private void updateBBRState(PathCongestionInfo info, int bytesSent) {
        long now = System.currentTimeMillis();
        if (now - info.lastProbeTime > info.bbrProbeInterval) {
            // 进入探测阶段
            info.bbrProbing = true;
            info.bbrProbeUp = !info.bbrProbeUp;
            info.lastProbeTime = now;
        }

        if (info.bbrProbing) {
            if (info.bbrProbeUp) {
                info.bbrCwnd *= (1 + GAMMA);
            } else {
                info.bbrCwnd *= (1 - GAMMA);
            }
        }
    }

    private void updateBBR(PathCongestionInfo info) {
        long now = System.currentTimeMillis();
        long rtt = now - info.lastSendTime;

        // 更新最小 RTT
        if (rtt < info.minRtt || now - info.minRttTime > MIN_RTT_WINDOW_MS) {
            info.minRtt = rtt;
            info.minRttTime = now;
        }

        // 更新带宽估计
        double deliveryRate = info.bytesAcked / (double) rtt;
        if (deliveryRate > info.maxBandwidth) {
            info.maxBandwidth = deliveryRate;
            info.bbrCwnd = info.maxBandwidth * info.minRtt;
        }
    }

    private void handleBBRLoss(PathCongestionInfo info) {
        info.bbrProbing = false;
        info.bbrCwnd = Math.max(info.bbrCwnd * 0.8, info.minRtt * info.maxBandwidth);
    }

    private void updateCUBIC(PathCongestionInfo info) {
        long now = System.currentTimeMillis();
        double t = (now - info.lastLossTime) / 1000.0; // 转换为秒
        double target = info.cubicC + Math.pow(t - info.cubicK, 3);

        if (info.cubicCwnd < info.cubicTarget) {
            // 慢启动
            info.cubicCwnd += 1.0;
        } else {
            // CUBIC 增长
            info.cubicCwnd = target;
        }
    }

    private void handleCUBICLoss(PathCongestionInfo info) {
        info.lastLossTime = System.currentTimeMillis();
        info.cubicTarget = info.cubicCwnd * 0.8;
        info.cubicCwnd = info.cubicTarget;
        info.cubicK = Math.cbrt(info.cubicTarget * (1 - 0.8) / info.cubicC);
    }

    /**
     * 内部每条子流的拥塞状态
     */
    private static class PathCongestionInfo {
        // 基本状态
        double cwnd;
        double ssthresh;
        private int inflightPackets = 0;
        private long bytesSent = 0;
        private long bytesAcked = 0;
        private long lastSendTime = 0;
        private long lastAckTime = 0;

        // BBR 状态
        double bbrCwnd = INITIAL_CWND;
        double maxBandwidth = 0;
        long minRtt = Long.MAX_VALUE;
        long minRttTime = 0;
        boolean bbrProbing = false;
        boolean bbrProbeUp = true;
        long lastProbeTime = 0;
        long bbrProbeInterval = 1000; // 1秒

        // CUBIC 状态
        double cubicCwnd = INITIAL_CWND;
        double cubicTarget = INITIAL_CWND;
        double cubicK = 0;
        double cubicC = 0.4;
        long lastLossTime = 0;

        // Jacobson RTT 平滑
        double srtt = 0;
        double rttvar = 0;

        PathCongestionInfo(double cwnd, double ssthresh) {
            this.cwnd = cwnd;
            this.ssthresh = ssthresh;
            this.lastSendTime = System.currentTimeMillis();
            this.lastLossTime = System.currentTimeMillis();
        }

        synchronized void updateRtt(long newRtt) {
            if (srtt == 0) {
                srtt = newRtt;
                rttvar = newRtt / 2.0;
            } else {
                double absErr = Math.abs(srtt - newRtt);
                rttvar = (1 - BETA) * rttvar + BETA * absErr;
                srtt = (1 - ALPHA) * srtt + ALPHA * newRtt;
            }
        }

        synchronized void packetSent() {
            inflightPackets++;
            lastSendTime = System.currentTimeMillis();
        }

        synchronized void packetAcked() {
            if (inflightPackets > 0) {
                inflightPackets--;
            }
            lastAckTime = System.currentTimeMillis();
        }

        synchronized int getInflight() {
            return inflightPackets;
        }

        synchronized void resetInflight() {
            inflightPackets = 0;
        }

        synchronized void addBytesSent(int bytes) {
            this.bytesSent += bytes;
        }

        synchronized void addBytesAcked(int bytes) {
            this.bytesAcked += bytes;
        }

        synchronized long getBytesSent() {
            return bytesSent;
        }

        synchronized long getBytesAcked() {
            return bytesAcked;
        }
    }
}
